import asyncio
import os
import json
from pathlib import Path
from urllib.parse import unquote
from io import BytesIO
from PIL import Image
from PIL.ExifTags import TAGS
from aiohttp import web


try:
    import aiofiles
    from aiofiles import os as aio_os  # 异步os操作
except ImportError:
    raise ImportError("请安装aiofiles库: pip install aiofiles")

class ImageService:
    default_node_field_map = {
        "CLIPTextEncode": "text",
        "WeiLinPromptUI": "positive",
        "WeiLinPromptUIWithoutLora": "positive",
        "MarvelousText": "prompt_text",
        "MarvelousBox": "prompt_text"
    }

    @classmethod
    async def read_metadata(cls, request: web.Request) -> web.Response:
        try:
            data = await request.post()
            file_field = data.get('file')
            if not file_field:
                return web.json_response({
                    "success": False,
                    "error": "缺少文件参数"
                }, status=400)

            user_map_str = data.get('map', '{}')
            try:
                user_node_field_map = json.loads(user_map_str)
                if not isinstance(user_node_field_map, dict):
                    raise ValueError("映射配置必须是对象类型")
            except json.JSONDecodeError:
                return web.json_response({
                    "success": False,
                    "error": "映射配置格式错误"
                }, status=400)

            merged_map = {**cls.default_node_field_map, **user_node_field_map}

            image_data = file_field.file.read()
            with Image.open(BytesIO(image_data)) as img:
                metadata = cls.extract_image_metadata(img)

            comfy_prompt = metadata.get('prompt')
            if comfy_prompt:
                parsed_result = cls.parse_comfyui_metadata(comfy_prompt, merged_map)
                if parsed_result['success']:
                    return web.json_response({
                        "success": True,
                        "metadata": {
                            "type": "comfyui",
                            "prompt": parsed_result['prompt'],
                            "positive": parsed_result['positive'],
                            "negative": parsed_result['negative'],
                            "rawPrompt": comfy_prompt
                        }
                    })
                else:
                    return web.json_response({
                        "success": True,
                        "metadata": {
                            "type": "comfyui_raw",
                            "prompt": f"解析失败: {parsed_result.get('error', '未知错误')}",
                            "rawPrompt": comfy_prompt
                        }
                    })

            sd_metadata = metadata.get('parameters')
            if sd_metadata:
                return web.json_response({
                    "success": True,
                    "metadata": {
                        "type": "sdwebui",
                        "prompt": sd_metadata.strip(),
                        "rawPrompt": sd_metadata.strip()
                    }
                })

            a1111_metadata = metadata.get('UserComment')
            if a1111_metadata and 'Steps:' in a1111_metadata:
                return web.json_response({
                    "success": True,
                    "metadata": {
                        "type": "a1111",
                        "prompt": a1111_metadata.strip(),
                        "rawPrompt": a1111_metadata.strip()
                    }
                })

            return web.json_response({
                "success": True,
                "metadata": {
                    "type": "unknown",
                    "prompt": "未找到可识别的元数据",
                    "rawPrompt": ""
                }
            })

        except Exception as e:
            return web.json_response(
                {"success": False, "error": str(e)},
                status=500
            )

    @staticmethod
    def extract_image_metadata(img: Image.Image) -> dict:
        """提取图片中的所有元数据（text和exif）"""
        metadata = {}

        if hasattr(img, 'text'):
            metadata.update(img.text)

        exif_data = img.getexif()
        if exif_data:
            for tag_id, value in exif_data.items():
                # 修复TAGS引用错误
                tag_name = TAGS.get(tag_id, tag_id)  # 现在TAGS已正确导入
                if isinstance(value, bytes):
                    try:
                        value = value.decode('utf-8', errors='replace')
                    except UnicodeDecodeError:
                        continue
                metadata[tag_name] = value

        return metadata

    @staticmethod
    def preprocess_json_string(json_str: str) -> str:
        return json_str.replace(': NaN', ': null')

    @classmethod
    def parse_comfyui_metadata(cls, comfy_prompt: str, node_field_map: dict) -> dict:
        try:
            processed_prompt = cls.preprocess_json_string(comfy_prompt)
            prompt_data = json.loads(processed_prompt)

            k_sampler = cls.find_ksampler_node(prompt_data)
            if not k_sampler:
                return {"success": False, "error": "未找到KSampler节点"}

            positive_node_id = k_sampler.get('inputs', {}).get('positive', [None])[0]
            negative_node_id = k_sampler.get('inputs', {}).get('negative', [None])[0]

            positive_text = cls.trace_prompt_chain(
                prompt_data, positive_node_id, node_field_map
            ) if positive_node_id else ""

            negative_text = cls.trace_prompt_chain(
                prompt_data, negative_node_id, node_field_map
            ) if negative_node_id else ""

            combined_prompt = positive_text
            if negative_text:
                combined_prompt += f"\nNegative prompt: {negative_text}"

            return {
                "success": True,
                "prompt": combined_prompt,
                "positive": positive_text,
                "negative": negative_text
            }

        except json.JSONDecodeError as e:
            return {"success": False, "error": f"JSON解析错误: {str(e)}"}
        except Exception as e:
            return {"success": False, "error": str(e)}

    @staticmethod
    def find_ksampler_node(prompt_data: dict) -> dict:
        for node in prompt_data.values():
            if isinstance(node, dict) and (node.get('class_type') == 'KSampler' or node.get('class_type') == 'KSamplerAdvanced' ):
                if 'positive' in node.get('inputs', {}) and 'negative' in node.get('inputs', {}):
                    return node
        return None

    # 将静态方法改为类方法，以便使用cls
    @classmethod
    def trace_prompt_chain(cls, prompt_data: dict, node_id: str, node_field_map: dict,
                           accumulated_text: str = "") -> str:
        """追踪提示词链，合并上下游文本"""
        if not node_id or node_id not in prompt_data:
            return accumulated_text

        node = prompt_data[node_id]
        class_type = node.get('class_type')
        inputs = node.get('inputs', {})

        current_text = ""
        if class_type in node_field_map:
            field_name = node_field_map[class_type]
            current_text = str(inputs.get(field_name, "")).strip()

        upstream_text = accumulated_text
        for input_value in inputs.values():
            if isinstance(input_value, list) and len(input_value) > 0:
                next_node_id = input_value[0] if isinstance(input_value[0], str) else None
                if next_node_id and next_node_id in prompt_data:
                    # 现在可以正确使用cls调用类方法
                    upstream_text = cls.trace_prompt_chain(
                        prompt_data, next_node_id, node_field_map, accumulated_text
                    )
                    break

        # 现在可以正确使用cls调用类方法
        return cls.merge_prompt_text(upstream_text, current_text)

    @staticmethod
    def merge_prompt_text(upstream: str, current: str) -> str:
        if not upstream:
            return current
        if not current:
            return upstream
        return f"{upstream}, {current}"

    @classmethod
    async def get_image(cls, request: web.Request) -> web.Response:
        path = request.query.get("path")
        if not path:
            return web.json_response(
                {"success": False, "error": "缺少路径参数"},
                status=400
            )
        try:
            decoded = unquote(path)
            pure_path = decoded.split("?")[0]
            if pure_path.startswith("file://"):
                # 统一替换为本地文件路径
                local_path = pure_path.replace("file://", "")
                # 处理Windows盘符后的斜杠（确保D:/或D:\正确保留）
                if local_path.startswith("/"):
                    local_path = local_path[1:]
                # 统一转换为Windows的反斜杠
                local_path = local_path.replace("/", "\\")
                full_path = Path(local_path).resolve()
            else:
                # 原有逻辑：处理相对路径
                full_path = (cls.get_target_path() / pure_path).resolve()

            # 4. 验证文件是否存在
            if not full_path.exists():
                return web.Response(status=404, text=f"文件不存在: {full_path}")
            if not full_path.is_file():
                return web.Response(status=400, text=f"目标不是文件: {full_path}")

            # 5. 验证文件格式
            suffix = full_path.suffix.lower()
            if suffix not in [".png", ".webp", ".jpg", ".jpeg"]:
                return web.Response(status=415, text=f"不支持的图片格式: {suffix}")

            return web.FileResponse(full_path)

        except Exception as e:
            return web.json_response(
                {"success": False, "error": str(e)},
                status=500
            )
    @classmethod
    async def get_metadata(cls, request: web.Request) -> web.Response:
        path = request.query.get("path")
        if not path:
            return web.json_response(
                {"success": False, "error": "缺少路径参数"},
                status=400
            )

        try:
            decoded = unquote(path)
            full_path = (cls.get_target_path() / decoded).resolve()

            if not full_path.exists():
                return web.Response(status=404, text="文件不存在")
            if not full_path.is_file():
                return web.Response(status=400, text="目标不是文件")
            if full_path.suffix.lower() != ".png":
                return web.Response(status=415, text="仅支持PNG文件")

            with Image.open(full_path) as img:
                metadata = img.text.get("parameters", "未找到元数据")

            return web.json_response({
                "success": True,
                "metadata": metadata,
                "file_path": str(full_path)
            })

        except Exception as e:
            return web.json_response(
                {
                    "success": False,
                    "error": f"元数据读取失败: {str(e)}"
                },
                status=500
            )

    @classmethod
    async def save_image_as_webp(cls, request: web.Request) -> web.Response:
        try:
            # 接收FormData中的文件和参数
            data = await request.post()

            # 获取图片文件
            file_field = data.get('file')
            if not file_field:
                return web.json_response({
                    "success": False,
                    "message": "未提供图片文件"
            })

            # 获取目标保存路径参数
            target_path_str = data.get('targetPath')
            if not target_path_str or not isinstance(target_path_str, str):
                return web.json_response({
                    "success": False,
                    "message": "无效的目标保存路径"
                })

            # 解码URL编码的路径
            decoded_path = unquote(target_path_str)

            # 处理不同格式的路径
            if decoded_path.startswith("file://"):
                local_path = decoded_path.replace("file://", "")
                if local_path.startswith("/"):
                    local_path = local_path[1:]
                local_path = local_path.replace("/", "\\")
                target_path = Path(local_path).resolve()
            else:
                # 处理绝对路径
                target_path = Path(decoded_path).resolve()

            # 确保目标路径是webp格式
            if target_path.suffix.lower() != '.webp':
                target_path = target_path.with_suffix('.webp')

            # 创建父目录（如果不存在）
            target_dir = target_path.parent
            if not await cls.file_exists(str(target_dir)):
                loop = request.app.loop
                await loop.run_in_executor(
                    None,
                    lambda: os.makedirs(target_dir, exist_ok=True, mode=0o755)
                )

            # 生成临时文件路径
            temp_file_name = f"{target_path.name}.tmp-{os.urandom(8).hex()}"
            temp_path = target_path.parent / temp_file_name

            try:
                # 读取上传的图片文件
                image_data = file_field.file.read()

                # 使用PIL处理图片并转换为WebP
                loop = request.app.loop

                # 在executor中处理图片转换（同步操作异步化）
                await loop.run_in_executor(None, cls._convert_image,
                                           image_data, str(temp_path))

                # 原子性重命名临时文件到目标路径
                await loop.run_in_executor(
                    None,
                    lambda: os.replace(str(temp_path), str(target_path))
                )

                return web.json_response({
                    "success": True,
                    "imagePath": str(target_path),
                    "message": "图片已成功保存为WebP格式"
                })

            except Exception as conversion_error:
                # 尝试清理临时文件
                if await cls.file_exists(str(temp_path)):
                    try:
                        await loop.run_in_executor(
                            None,
                            lambda: os.unlink(str(temp_path))
                        )
                    except Exception as cleanup_error:
                        print(f"清理临时文件失败: {str(cleanup_error)}")

                return web.json_response({
                    "success": False,
                    "message": f"图片转换失败: {str(conversion_error)}"
                })

        except Exception as main_error:
            return web.json_response({
                "success": False,
                "message": f"保存图片时出错: {str(main_error)}"
            })

    @staticmethod
    async def file_exists(file_path: Path | str) -> bool:
        """优化后的异步检查文件是否存在（兼容str和Path类型）"""
        try:
            # 先将输入转换为Path（无论原始类型是str还是Path）
            path = Path(file_path) if isinstance(file_path, str) else file_path
            return await asyncio.to_thread(path.exists)
        except Exception:
            # 发生异常时默认返回不存在（例如权限不足）
            return False
    @staticmethod
    def _convert_image(image_data: bytes, output_path: str) -> None:
        """同步图片转换方法，用于在executor中运行"""
        with Image.open(BytesIO(image_data)) as img:
            # 确保图片是RGB模式（处理透明通道）
            if img.mode in ('RGBA', 'LA'):
                background = Image.new(img.mode[:-1], img.size, (255, 255, 255))
                background.paste(img, img.split()[-1])
                img = background
            elif img.mode == 'P':
                img = img.convert('RGB')

            # 保存为WebP格式，质量90
            img.save(output_path, 'WebP', quality=90, lossless=False)